package net.contra.jmd.transformers.allatori;

import net.contra.jmd.transformers.Transformer;
import net.contra.jmd.util.GenericMethods;
import net.contra.jmd.util.LogHandler;
import net.contra.jmd.util.NonClassEntries;
import org.apache.bcel.classfile.ClassParser;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.*;

import java.io.File;
import java.lang.reflect.InvocationTargetException;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.regex.Pattern;

public class AllatoriTransformer implements Transformer {
    LogHandler logger = new LogHandler("AllatoriTransformer");
    Map<String, ClassGen> cgs = new HashMap<String, ClassGen>();
    Map<String, Class> loadedClasses = new HashMap<>();
    URLClassLoader urlClassLoader;
    ClassGen ALLATORI_CLASS;
    String JAR_NAME;
    boolean isStrong = false;
    private static Pattern pattern=Pattern.compile("");
    public AllatoriTransformer(String jarfile, boolean strong) throws Exception {
        logger.log("Allatori Deobfuscator");
        isStrong = strong;
        File jar = new File(jarfile);
        JAR_NAME = jarfile;
        JarFile jf = new JarFile(jar);
        Enumeration<JarEntry> entries = jf.entries();
        while (entries.hasMoreElements()) {
            JarEntry entry = entries.nextElement();
            if (entry == null) {
                break;
            }
            if (entry.getName().endsWith(".class")) {
                ClassGen cg = new ClassGen(new ClassParser(jf
                        .getInputStream(entry), entry.getName()).parse());
                if (isStringClass(cg) || isStringClassB(cg)) {
                    ALLATORI_CLASS = cg;
                    logger.debug("Allatori Class: " + ALLATORI_CLASS.getClassName());
                }
                cgs.put(cg.getClassName(), cg);
            } else {
                NonClassEntries.add(entry, jf.getInputStream(entry));
            }
        }
        jf.close();
    }

    private boolean isStringClass(ClassGen cg) {
        if (cg.getMethods().length == 2 && cg.getMethods()[0].isStatic()
                && cg.getMethods()[1].isStatic()) {
            if (cg.getMethods()[0].getReturnType().toString().equals(
                    "java.lang.String")
                    && cg.getMethods()[1].getReturnType().toString().equals(
                    "java.lang.String")) {
                return true;
            }
        }
        return false;
    }

    private boolean isStringClassB(ClassGen cg) {
        if (cg.getMethods().length == 1 && cg.getMethods()[0].isStatic()) {
            if (cg.getMethods()[0].getReturnType().toString().equals("java.lang.String")) {
                return true;
            }
        }
        return false;
    }
    public void transform(String jarfile) throws TargetLostException, MalformedURLException, ClassNotFoundException, IllegalAccessException, InstantiationException, NoSuchMethodException, InvocationTargetException {
        logger.log("Starting load class");
        String dir = "D:/workspace/jeesite4/web/target/web/WEB-INF/lib/";
        String pathUrl = "file:" + dir;
        File file = new File(dir);
        String[] list = file.list();
        URL[] urls = new URL[list.length];
        for (int i = 0; i < list.length; i++) {
            urls[i] = new URL(pathUrl + list[i]);
        }
        urlClassLoader = new URLClassLoader(urls, Thread.currentThread().getContextClassLoader());
        logger.log("Starting Encrypted String Removal...");
        replaceStrings();
        logger.log("Deobfuscation Finished!");
        logger.log("Starting remove garbage code");
//        removeGarbageCode();
        logger.log("Finished remove garbage code");
        GenericMethods.dumpJar(JAR_NAME, cgs.values());
        logger.log("Operation Completed.");
    }

    private void removeGarbageCode() throws ClassNotFoundException {
        for (ClassGen cg : cgs.values()) {
            int replaced = 0;
            for (Method method : cg.getMethods()) {
                MethodGen mg = new MethodGen(method, cg.getClassName(), cg.getConstantPool());
                InstructionList list = mg.getInstructionList();
                if (list == null) {
                    continue;
                }
                InstructionHandle[] handles = list.getInstructionHandles();
                for (int i=0;i<handles.length;i++){
                    InstructionHandle handle = handles[i];
                    //去除无用的汇编指令
                }
                mg.setInstructionList(list);
                mg.setMaxLocals();
                mg.setMaxStack();


                cg.replaceMethod(method, mg.getMethod());
            }
            if (replaced > 0) {
                logger.debug("decrypted " + replaced + " strings in class " + cg.getClassName());
            }
        }
    }

    public void replaceStrings() throws ClassNotFoundException {
        int total=0;
        int totalException=0;

        for (ClassGen cg : cgs.values()) {
            int replaced = 0;
            for (Method method : cg.getMethods()) {
                MethodGen mg = new MethodGen(method, cg.getClassName(), cg.getConstantPool());
                InstructionList list = mg.getInstructionList();
                if (list == null) {
                    continue;
                }
                InstructionHandle[] handles = list.getInstructionHandles();
                for (int i = 1; i < handles.length; i++) {
                    Instruction invokeString = handles[i - 1].getInstruction();
                    Instruction invokeStaticMethod = handles[i].getInstruction();
                    if (invokeStaticMethod instanceof INVOKESTATIC && invokeString instanceof LDC) {
                        //需要反混淆的条件，1输入参数是固定字符串，2有返回值
                        LDC encryptedLDC = (LDC) invokeString;
                        INVOKESTATIC methodCall = (INVOKESTATIC) invokeStaticMethod;
                        String methodName = methodCall.getMethodName(cg.getConstantPool());
                        Type[] argumentTypes = methodCall.getArgumentTypes(cg.getConstantPool());
                        if (!methodName.equals("ALLATORIxDEMO") ) {
                            continue;
                        }
                        if (argumentTypes.length!=1||!argumentTypes[0].equals(Type.STRING )){
                            continue;
                        }
                        String className = methodCall.getClassName(cg.getConstantPool());
                        Class clazz;
                        if (loadedClasses.containsKey(className)) {
                            clazz = loadedClasses.get(className);
                        } else {
                            clazz = urlClassLoader.loadClass(className);
                            loadedClasses.put(className, clazz);
                        }


                        String encryptedString = encryptedLDC.getValue(cg.getConstantPool()).toString();
                        String decryptedString = null;
                        try {
//                            if (cg.getClassName().equals("com.jeesite.common.shiro.l.a")&&method.getName().equals("ALLATORIxDEMO")&&method.getArgumentTypes()!=null&&method.getArgumentTypes().length==1){
//                                System.out.println();
//                            }
                            java.lang.reflect.Method[] declaredMethods = clazz.getDeclaredMethods();
                            for (java.lang.reflect.Method method1:declaredMethods){
                                if (method1.getReturnType().getName().equals("java.lang.String")&&method1.getName().equals("ALLATORIxDEMO") ){
                                    Class<?>[] parameterTypes = method1.getParameterTypes();
                                    if (parameterTypes !=null&&parameterTypes.length==1&&parameterTypes[0].getName().equals("java.lang.String")){
                                        decryptedString = (String) method1.invoke(null, encryptedString);
                                        logger.debug(encryptedString + " -> " + decryptedString + " in " + cg.getClassName() + "." + method.getName());
                                        int stringRef = cg.getConstantPool().addString(decryptedString);
                                        LDC lc = new LDC(stringRef);
                                        NOP nop = new NOP();
                                        handles[i].setInstruction(lc);
                                        handles[i - 1].setInstruction(nop);
                                        replaced++;
                                    }
                                }
                            }
                        } catch (Exception e) {
                            e.printStackTrace();
                            totalException++;
                            logger.error("error replace method from class=" + cg.getClassName() + " method =" + method.getName() + " targetClass=" + className + " targetMethod=" + methodName + " cause=" + e.getMessage());
                        }
                    }
                }
                mg.setInstructionList(list);
                mg.setMaxLocals();
                mg.setMaxStack();
                cg.replaceMethod(method, mg.getMethod());
            }
            if (replaced > 0) {
                logger.debug("decrypted " + replaced + " strings in class " + cg.getClassName());
                total+=replaced;
            }
        }
        logger.debug("恭喜 所有字符串解密完成，共解密 " + total + " 个字符串,异常个数为"+totalException+"个");
    }
}
